package main.utils

case class TreeNode[T](value: T, children: List[TreeNode[T]] = Nil)

// Tree implementation (each node's value is unique)
// TODO: Support for common values
class Tree[T](val root: TreeNode[T]) {
  def this(rootValue: T) = {
    this(TreeNode(rootValue))
  }

  def findNode(value: T): Option[TreeNode[T]] = {
    def findRecursively(node: TreeNode[T]): Option[TreeNode[T]] =
      if (node.value == value) Some(node)
      else node.children.map(findRecursively).find(_.isDefined).flatten

    findRecursively(root)
  }

  def addChild(parentData: T, childData: T): Tree[T] = {
    val newRoot = addNode(root, parentData, TreeNode(childData))
    new Tree(newRoot)
  }

  def addNode(
      node: TreeNode[T],
      parentData: T,
      child: TreeNode[T]
  ): TreeNode[T] = {
    if (node.value == parentData)
      node.copy(children = node.children :+ child)
    else
      node.copy(children = node.children.map(addNode(_, parentData, child)))

  }

  def removeChild(parentData: T, childData: T): Tree[T] = {
    val newRoot = removeNode(root, parentData, TreeNode(childData))
    new Tree(newRoot)
  }

  def removeNode(
      node: TreeNode[T],
      parentData: T,
      child: TreeNode[T]
  ): TreeNode[T] = {
    if (node.value == parentData)
      node.copy(children = node.children.filterNot(_ == child))
    else
      node.copy(children = node.children.map(removeNode(_, parentData, child)))
  }

  def printTree(): Unit = {
    def printNode(node: TreeNode[T], depth: Int): Unit = {
      println("  " * depth + node.value.toString)
      node.children.foreach(printNode(_, depth + 1))
    }

    printNode(root, 0)
  }

}
